package mw

import (
	"context"
	"hertz/demo/biz/model"
	"hertz/demo/biz/util"
	"time"

	"github.com/cloudwego/hertz/pkg/app"
	"github.com/cloudwego/hertz/pkg/common/utils"
	"github.com/cloudwego/hertz/pkg/protocol/consts"
	"github.com/hertz-contrib/jwt"
)

// user in jwt token
type JWTUser struct {
	Id    uint32 `json:"id"`
	Name  string `json:"name"`
	Admin bool   `json:"admin"`
}

type LoginReq struct {
	UserName string `json:"username"`
	PassWord string `json:"password"`
}

const IDENTITY_KEY string = "id"
const NAME_KEY string = "name"
const ADMIN_KEY string = "admin"

func BuildAuthMiddleware(secretKey []byte, timeout time.Duration) (*jwt.HertzJWTMiddleware, error) {
	return jwt.New(&jwt.HertzJWTMiddleware{
		Realm:        "label system",
		Key:          secretKey,
		Timeout:      timeout,
		MaxRefresh:   timeout,
		SendCookie:   true,
		TokenLookup:  "header: Authorization, cookie: jwt",
		CookieName:   "jwt",
		SecureCookie: true,
		IdentityKey:  IDENTITY_KEY,
		PayloadFunc: func(data interface{}) jwt.MapClaims {
			if v, ok := data.(*JWTUser); ok {
				return jwt.MapClaims{
					IDENTITY_KEY: v.Id,
					NAME_KEY:     v.Name,
					ADMIN_KEY:    v.Admin,
				}
			}
			return jwt.MapClaims{}
		},
		IdentityHandler: func(ctx context.Context, c *app.RequestContext) interface{} {
			claims := jwt.ExtractClaims(ctx, c)
			return &JWTUser{
				Id:    uint32(claims[IDENTITY_KEY].(float64)),
				Name:  claims[NAME_KEY].(string),
				Admin: claims[ADMIN_KEY].(bool),
			}
		},
		LoginResponse: func(ctx context.Context, c *app.RequestContext, code int, token string, expire time.Time) {
			c.JSON(code, utils.H{
				"token":  token,
				"expire": expire.Format(time.RFC3339),
			})
		},
		RefreshResponse: func(ctx context.Context, c *app.RequestContext, code int, token string, expire time.Time) {
			c.JSON(code, utils.H{
				"token":  token,
				"expire": expire.Format(time.RFC3339),
			})
		},

		Authenticator: func(ctx context.Context, c *app.RequestContext) (interface{}, error) {
			var req LoginReq
			if err := c.BindAndValidate(&req); err != nil {
				return nil, jwt.ErrMissingLoginValues
			}

			var users []model.User
			GetDB(c).Table("users").Where("username = ?", req.UserName).Find(&users)
			if len(users) == 0 {
				return nil, jwt.ErrFailedAuthentication
			}

			match := util.CheckPasswordHash(req.PassWord, users[0].Password)
			if match && !users[0].Disable {
				return &JWTUser{
					Id:    uint32(users[0].ID),
					Name:  users[0].Username,
					Admin: users[0].IsAdmin,
				}, nil
			} else {
				return nil, jwt.ErrFailedAuthentication
			}
		},
		Authorizator: func(data interface{}, ctx context.Context, c *app.RequestContext) bool {
			if user, ok := data.(*JWTUser); ok {
				enable, err := model.IsUserEnable(GetDB(c), uint(user.Id))
				if err != nil {
					return false
				}
				return enable
			}

			return false
		},
		Unauthorized: func(ctx context.Context, c *app.RequestContext, code int, message string) {
			c.JSON(code, map[string]interface{}{
				"message": message,
			})
		},
	})

}

// 中间件, 要求用户必须为管理员
func RequireAdmin() app.HandlerFunc {
	return func(ctx context.Context, c *app.RequestContext) {
		u, exist := c.Get(IDENTITY_KEY)
		if !exist {
			c.JSON(consts.StatusUnauthorized, utils.H{
				"message": "Unauthorized",
			})
			c.Abort()
			return
		}
		user, ok := u.(*JWTUser)
		if !ok {
			c.JSON(consts.StatusUnauthorized, utils.H{
				"message": "Unauthorized",
			})
			c.Abort()
			return
		}

		if user.Admin {
			c.Next(ctx)
			return
		} else {
			c.JSON(consts.StatusUnauthorized, utils.H{
				"message": "Unauthorized",
			})
			c.Abort()
			return
		}
	}
}

func GetUser(c *app.RequestContext) *JWTUser {
	u, exist := c.Get(IDENTITY_KEY)
	if !exist {
		return nil
	}
	user, ok := u.(*JWTUser)
	if !ok {
		panic("invalid request context")
	}
	return user
}
